from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, Dropout
from keras.layers import Flatten, Dense
from keras import regularizers

def classification (batch_size = 20, epochs = 20, img_width = 150,
                    img_height = 150, 
                    num_train_samples = 2000, num_validation_samples=800):
    train_data_dir = 'Datasets/DogsCats/train'
    validation_data_dir = 'Datasets/DogsCats/validation'

    model = Sequential()
    model.add(Conv2D(32, (3, 3), activation='relu',
              input_shape=(img_width, img_height, 3)))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(48, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    #model.add(Dropout(0.4))      # Dropout
    model.add(Flatten())
    model.add(Dense(64, activation='relu',
              kernel_regularizer=regularizers.l2(0.01)))
    model.add(Dense(1, activation='sigmoid',  ))

    model.compile(loss='binary_crossentropy',
              optimizer='rmsprop',
              metrics=['accuracy'])

    train_datagen = ImageDataGenerator(rescale=1. / 255)
    test_datagen = ImageDataGenerator(rescale=1. / 255)

    train_generator = train_datagen.flow_from_directory(
        train_data_dir,
        target_size=(img_width, img_height),
        batch_size=batch_size,
        class_mode='binary')

    validation_generator = test_datagen.flow_from_directory(
        validation_data_dir,
        target_size=(img_width, img_height),
        batch_size=batch_size,
        class_mode='binary')

    history = model.fit_generator(
        train_generator,
        steps_per_epoch=num_train_samples // batch_size,
        epochs=epochs,
        validation_data=validation_generator,
        validation_steps=num_validation_samples // batch_size,
        verbose =0)
    
    return history

history = classification(num_train_samples = 4000)
plot(history)
